#include "Arduino.h"
#include <WiFi.h>
#include "./src/wifi.h"
#include "./src/inmp441.h"
#include "./src/tcp_client.h"
#include "./src/codec.h"
#include "./src/ws_client.h"

static int isOK = 0;
static uint64_t raw_length = 0;
static uint64_t total_encoded_length = 0;
static int total_count = 0;

#define READ_BUFF_SIZE FRAME_SAMPLES * 2

static int16_t i2s_read_buff[READ_BUFF_SIZE];
static uint8_t encoded_opus_buff[FRAME_BYTES * 4];

void wsEncodeTask(void *param) {
    Serial.println("WebSocket opus start");
    
    while (true) {
        // 定期检查连接健康状态
        ws_check_connection_health();
        
        size_t len = inmp441_listen(i2s_read_buff, READ_BUFF_SIZE * sizeof(int16_t));
        raw_length += len;

        int16_t *p = i2s_read_buff;
        size_t samples_left = len / sizeof(int16_t);
        while (samples_left >= FRAME_SAMPLES) {
            int encoded_size = encode(p, encoded_opus_buff, sizeof(encoded_opus_buff));
            if (encoded_size < 0) {
                Serial.println("Encoding failed");
                break;
            }
            if (encoded_size > 0) {
                ws_send_bytes(encoded_opus_buff, encoded_size);
                if (ws_is_connected()) {
                    total_count++;
                    total_encoded_length += encoded_size;
                }
            }
            p += FRAME_SAMPLES;
            samples_left -= FRAME_SAMPLES;
        }

        vTaskDelay(1 / portTICK_PERIOD_MS);
    }
    // vTaskDelete(NULL);
}

void encodeTask(void *param)
{
  isOK = tcp_connect();
  while (true)
  {
    if (isOK == 0)
    {
      return;
    }

    size_t len = inmp441_listen(i2s_read_buff, READ_BUFF_SIZE * sizeof(int16_t));
    int current_encode_size = 0;

    int16_t *p = i2s_read_buff;
    size_t samples_left = len / sizeof(int16_t);

    while (samples_left >= FRAME_SAMPLES)
    {
      int encoded_size = encode(p, encoded_opus_buff, sizeof(encoded_opus_buff));
      if (encoded_size < 0)
      {
        Serial.println("Encoding failed");
        break;
      }
      if (encoded_size > 0)
      {
        tcp_send_bytes(encoded_opus_buff, encoded_size);
        total_count++;
        raw_length += FRAME_BYTES;
        total_encoded_length += encoded_size;
        current_encode_size += encoded_size;
      }
      p += FRAME_SAMPLES;
      samples_left -= FRAME_SAMPLES;
    }
    Serial.printf("finish encoded bytes %d, raw bytes: %d\n", current_encode_size, len);
  }
}

void setup()
{
  Serial.begin(115200);
  randomSeed(analogRead(0)); // 初始化随机数种子
  connect_wifi();
  WiFi.setSleep(false);
  
  // 启动时运行网络诊断
  Serial.println("Running initial network diagnostics...");
  ws_network_diagnostics();
  
  inmp441_init();
  initCodec();

  // xTaskCreatePinnedToCore(
  //     encodeTask,
  //     "encodeTask",
  //     8192 * 4,
  //     NULL,
  //     1,
  //     NULL,
  //     0);

  xTaskCreatePinnedToCore(
      wsEncodeTask,
      "wsEncodeTask",
      8192 * 4,
      NULL,
      3,
      NULL,
      0);
}

void loop()
{
  static unsigned long lastLog = 0;
  if (millis() - lastLog > 3000)
  {
    if (ws_is_connected()) {
      Serial.printf("WS: connected | frames: %d | raw: %llu bytes | encoded: %llu bytes\n", 
                    total_count, (unsigned long long)raw_length, (unsigned long long)(total_encoded_length));
    } else {
      int retryCount = ws_get_retry_count();
      Serial.printf("WS: retry(%d) | encoded: %d | raw: %llu bytes\n", 
                    retryCount, total_count, (unsigned long long)raw_length);
    }
    lastLog = millis();
  }
}
